package arctic;

import core.HiddenLayer;
import core.Phone;
import core.Speaker;

import java.io.*;
import java.util.*;

/**
 * Created with IntelliJ IDEA.
 * User: lahiru
 * Date: 5/27/14
 * Time: 1:05 PM
 * To change this template use File | Settings | File Templates.
 */
public class CollectStats {

    private static final String pAlignments = "/media/Bunty/Academic/Research/arctic-speaker-effects/phone-alignments/phones_frame.ark";
    private static final String hiddenActivationDir = "/media/Bunty/Academic/Research/arctic/dnn/activations/";
    private static final String phoneSets = "/media/Bunty/Academic/Research/arctic-speaker-effects/sets.int";
    private static final String saveObjects = "/media/Bunty/Academic/Research/arctic/dnn/objects/";
    private static final Map<String, Integer> phoneMapping = new HashMap<String, Integer>();
    private static final String spk="awb";
    private static final int layers = 6;
    private static final int numberOfHiddenUnits=512;

    public static void main(String[] args) {



        //create the phoneMap
        try {
            BufferedReader mapBr = new BufferedReader(new FileReader(phoneSets));
            String line;
            int value = 1;
            while ((line = mapBr.readLine()) != null) {
                String phoneIds[] = line.split(" ");
                for (String phn : phoneIds) {
                    phoneMapping.put(phn, value);
                }
                value = value + 1;
            }

            mapBr.close();
        } catch (IOException e) {
            e.printStackTrace();  //To change body of catch statement use File | Settings | File Templates.
        }


        for (int i = 1; i <= layers; i++) {

            System.out.println(" Start processing the layer " + i);
            HiddenLayer hiddenLayer = new HiddenLayer(i);
            int finalFrameCount = 0;
            Speaker speaker = new Speaker();
            speaker.setId(spk);


            String ha = hiddenActivationDir + "h" + i + "/activations.ark";

            System.out.println(" Reading the activation file  " + ha + " for the speaker " + spk);
            System.out.println(" phone labels are in  " + pAlignments + " for the speaker " + spk);

            //Read phone-alignments line by line, file contains a line per utterance.

            try {
                BufferedReader phonesBr = new BufferedReader(new FileReader(pAlignments));
                String line;
                List<Integer> phoneList = new ArrayList<Integer>();


                //Now for each utterance read the hidden activations frame by frame.
                BufferedReader haBr = new BufferedReader(new FileReader(ha));
                String line1;

                int utteranceCount = 0;
                int frameCount = 0;
                String frameUtterance = "";
                String utterance = "";


                //read a line from the phone-alignments for each frame in the utterance
                while ((line = phonesBr.readLine()) != null) {
                    String phoneIds[] = line.split(" ");
                    utterance = phoneIds[0];
//                                System.out.println("The utterence is " +utterance);
                    phoneList = getMappedPhones(Arrays.asList(Arrays.copyOfRange(phoneIds, 1, phoneIds.length)), phoneMapping);
//                                System.out.println("The size of the phone list is  " + phoneList.size());


                    while ((line1 = haBr.readLine()) != null) {

                        String hArray[] = line1.trim().split(" ");

                        if (hArray.length == 3) {  //starting of an utterance
                            utteranceCount = utteranceCount + 1;
                            frameCount = 0;
                            frameUtterance = hArray[0];

                            String tempArray[] = null;
                            String line2 = "";

                            if (!frameUtterance.equalsIgnoreCase(utterance)) {
                                System.err.println("The frame utterance " + frameUtterance);
                                System.err.println("The phone utterance " + utterance);

                                do {
                                    line2 = haBr.readLine();
                                    tempArray = line2.trim().split(" ");

                                } while (tempArray.length != 3);

                                System.err.println("The new frame utterance is " + tempArray[0]);

                            }

//                            //read a line from the phone-alignments for each frame in the utterance
//                            if ((line = phonesBr.readLine()) != null){
//                                String phoneIds[] = line.split(" ");
//                                utterance = phoneIds[0];
//                                phoneList = getMappedPhones(Arrays.asList(Arrays.copyOfRange(phoneIds,1,phoneIds.length)), phoneMapping);
//                            } else{
//                                System.err.println("There is an error in files ");
//                                System.exit(1);
//                            }

                        } else if (hArray.length == numberOfHiddenUnits) {  //intermediate frame
//                            try {
                            speaker.addFrameStats(phoneList.get(frameCount), hArray);
//                            } catch (IndexOutOfBoundsException e) {
//                                System.err.println("The frame count is " + frameCount);
//                                System.err.println("The frame utterance is " + frameUtterance);
//                                System.err.println("The phone alignment utterance is " + utterance);
//                                System.err.println(" the error is " + e.getMessage());
//                            }
                            frameCount = frameCount + 1;
                        } else if (hArray.length == numberOfHiddenUnits +1) {   //last frame of the utterance
                            speaker.addFrameStats(phoneList.get(frameCount), Arrays.copyOfRange(hArray, 0, hArray.length - 1));
                            frameCount = frameCount + 1;
                            finalFrameCount = finalFrameCount + frameCount;
                            break;
                        } else {
                            System.err.println("This is not possible " + hArray.length);
                        }
                    }
                }
                haBr.close();
                System.out.println(" The utterance count " + utteranceCount);
                System.out.println(" The finalFrameCount " + finalFrameCount);

                phonesBr.close();
            } catch (IOException e) {
                e.printStackTrace();  //To change body of catch statement use File | Settings | File Templates.
            }


            //Printout speaker data,
            int totalPhoneCount = 0;
            Map<Integer, Phone> phones = speaker.getPhones();
            for (int pid : phones.keySet()) {
                Phone phone = phones.get(pid);
                //phone.averageVector(); //TODO should only do once when finished
                int count = phone.getCount();
                totalPhoneCount = totalPhoneCount + count;
//                    System.out.println("phone id " + pid + " count " + count);
            }

            if (totalPhoneCount != finalFrameCount) {
                System.err.println(" There is a mismatch of frames for the " + speaker.getId());
            } else {
                System.out.println("Frame counts are matching for the  " + speaker.getId());
            }


            hiddenLayer.getSpeakerList().add(speaker); //add the current speaker to the hidden layer



            //save the hidden layer data
            String location = null;
            try {
                location = saveObjects + "h" + i + ".ser";
                FileOutputStream fileOut =
                        new FileOutputStream(location);
                ObjectOutputStream out = new ObjectOutputStream(fileOut);
                out.writeObject(hiddenLayer);
                out.close();
                fileOut.close();
            } catch (IOException e) {
                e.printStackTrace();  //To change body of catch statement use File | Settings | File Templates.
            }
            System.out.println("Serialized data is saved in " + location + " for " + hiddenLayer.getSpeakerList().size() + " speakers ");


//            System.out.println(" Finished processing the layer " + i);
        }


        //(ArrayList<String>)
    }

    private static List<Integer> getMappedPhones(List<String> originalPhones, Map<String, Integer> phoneMapping) {
        List<Integer> phoneList = new ArrayList<Integer>(originalPhones.size());
        for (String phn : originalPhones) {
            phoneList.add(phoneMapping.get(phn));
        }
        return phoneList;
    }
}
